from zipline.api import order, symbol, get_datetime, record, set_commission, set_slippage
from zipline.finance.commission import PerTrade
from zipline.finance.slippage import FixedSlippage

def initialize(context):
    # 设置策略参数
    context.atr_period = 14
    context.atr_multiplier = 1.5
    context.exit_time = 15  # 平仓时间(小时)
    
    # 设置股票
    context.asset = symbol('AAPL')
    
    # 设置佣金和滑点
    set_commission(PerTrade(0.001))  # 0.1% 佣金
    set_slippage(FixedSlippage(0.01))  # 1% 滑点
    
    # 初始化变量
    context.prev_close = None
    context.atr = None

def handle_data(context, data):
    current_time = get_datetime()
    
    # 更新前日收盘价(每日开盘时)
    if current_time.hour == 9 and current_time.minute == 30:
        # 获取昨日收盘价
        context.prev_close = data.history(context.asset, 'close', 1, '1d')[-1]
        
        # 计算ATR需要足够的历史数据
        if len(data.history(context.asset, ['high','low','close'], context.atr_period+1, '1d')) >= context.atr_period:
            highs = data.history(context.asset, 'high', context.atr_period, '1d')
            lows = data.history(context.asset, 'low', context.atr_period, '1d')
            closes = data.history(context.asset, 'close', context.atr_period, '1d')
            
            # 手动计算ATR
            trs = []
            for i in range(1, len(highs)):
                high = highs[i]
                low = lows[i]
                prev_close = closes[i-1]
                tr = max(high - low, abs(high - prev_close), abs(low - prev_close))
                trs.append(tr)
                
            context.atr = sum(trs) / len(trs)
    
    if context.prev_close is None or context.atr is None:
        return
        
    # 计算突破阈值
    upper_band = context.prev_close + context.atr * context.atr_multiplier
    lower_band = context.prev_close - context.atr * context.atr_multiplier
    
    # 获取当前价格
    current_high = data.current(context.asset, 'high')
    current_low = data.current(context.asset, 'low')
    
    # 交易信号
    if not context.portfolio.positions[context.asset].amount:
        if current_high >= upper_band:
            order(context.asset, 100)  # 向上突破买入
    elif current_low <= lower_band:
        order(context.asset, -100)  # 向下突破卖出
            
    # 尾盘平仓
    if context.portfolio.positions[context.asset].amount and current_time.hour >= context.exit_time:
        order(context.asset, -context.portfolio.positions[context.asset].amount)